from evalscope import TaskConfig, run_task
from evalscope.constants import EvalType

import argparse

def get_task(args):
    return TaskConfig(
        model=args.serve_id,   # 模型名称 (需要与部署时的模型名称一致)
        api_url=f'http://{args.ip}:{args.port}/v1',  # 推理服务地址
        api_key='EMPTY',
        eval_type=EvalType.SERVICE,   # 评测类型，SERVICE表示评测推理服务
        datasets=[
            'math_500',
            'gsm8k',
            'mmlu',
            'gpqa',
            'aime24',
            'ceval',
        ],
        dataset_args={ # EvalScope内置支持，无需指定数据集ID
            'ceval': {'few_shot_num': 0 },
            'mmlu': {'few_shot_num': 0 },
            'gsm8k': {'few_shot_num': 0 },
            'math_500': {'few_shot_num': 0 } ,
            'gpqa': {'subset_list': ['gpqa_diamond'], 'few_shot_num': 0},
            'aime24': {'few_shot_num': 0},
        },
        eval_batch_size=args.bs,      # 发送请求的并发数
        generation_config={       # 模型推理配置
            'max_tokens': args.max_tokens,  # 最大生成token数，建议设置为较大值避免输出截断
            'temperature': 0.6,   # 采样温度 
            'top_p': 0.95,        # top-p采样
            'n': 1                # 每个请求产生的回复数量
        },
        stream=False,               # 是否使用流式请求，推荐设置为True防止请求超时
        timeout=360000
    )

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--serve_id", type=str, default="/home/weights/DeepSeek-R1-int4-tnv2")
    parser.add_argument("--ip", type=str, default="127.0.0.1")
    parser.add_argument("--port", type=int, default=8000)
    parser.add_argument("--bs", type=int, default=8)
    parser.add_argument("--max_tokens", type=int, default=1024*44)
    args = parser.parse_args()
    
    task_cfg = get_task(args)
    
    # while True:
    run_task(task_cfg=task_cfg)
